#include "mm.h"
#include "kheap.h"
#include "typedef.h"
#include "vga.h"
#include "string.h"
#include "kernel.h"
#include "asm.h"
#include "interrupts.h"

// this is the original pde as defined in cr1.s
extern u32 pde_orig;
// keeps track of pages we need to keep in memory
big_page_desc* big_page_desc_list;
little_page_desc* little_page_desc_list;

void page_fault();

// starts up the memory management system
void init_mm() {
  u32 i;
  // setup kernel pde to point to our tables
  kernel_dir = &pde_orig;
  cur_dir = kernel_dir;
  //kernel_dir->tables[0] = 0;
  // register page fault handler.
  register_isr(14, page_fault);
  // Now, enable paging!
  switch_page_directory(kernel_dir);
  init_heap(HEAP_START, HEAP_INIT_SIZE);
  init_phys_mm();
  init_virt_mm();
  // startup big page descriptor list. first element is never freed
  big_page_desc_list = (big_page_desc*)kmalloc(sizeof(big_page_desc));
  big_page_desc_list->addr = (u32)kernel_dir->tables[PDE_OFFSET(KERNEL_VIRT_START)] + KERNEL_VIRT_OFFSET;
  big_page_desc_list->next = big_page_desc_list;
  big_page_desc_list->prev = big_page_desc_list;
  big_page_desc_list->free = FALSE;
  big_page_desc_list->in_use = TRUE;
  big_page_desc_list->ref_count = 1;
  // same goes for the little list
  little_page_desc_list = (little_page_desc*)kmalloc(sizeof(little_page_desc));
  little_page_desc_list->addr = 0;
  little_page_desc_list->next = little_page_desc_list;
  little_page_desc_list->prev = little_page_desc_list;
}

// page fault handler. detects if fault was writing to a cow table, and 
// appropriatley perform the copy
void page_fault() {
  u32 addr = get_CR2();
  /*  kprintf("page fault at %x\n", addr);
      return;*/

  // this fault was a result of trying to write to read only page
  if(addr & (PG_PRES | PG_USER | PG_WRITE)) {
    // if this was a cow fault, find the table this one references and copy
    // the appropriate frame to a fresh one in this table
    if(addr & PG_COW) {
       kprintf("  Page Fault with flags at %p\n\0", addr);
       while(1);
    }
  }

  else {
    kprintf("  Page Fault at %p\n\0", addr);
    while(1);
  }
}

// creates a user for the specified range and maps in kernel space
void map_area(pde_t* dir, u32 start, u32 end, u32 flags) {
  u32 i, idx;
  pte_t* table;
  
  for(i = PG_ALIGN(start); i < PG_ALIGN(end); i += 0x1000) {
    if(dir->tables[PDE_OFFSET(i)] == 0) {
      table = (pte_t*)get_page();
      memset((u32)table, 0, sizeof(pte_t));
      dir->tables[PDE_OFFSET(i)] = (pte_t*)(((u32)table - KERNEL_VIRT_OFFSET) | flags);
    }
  
    idx = get_frame();
    table->pages[PTE_OFFSET(i)] = idx | flags;
  }
  
  return dir;
}  

// clones page directory and page tables. marks frames as copy on write
pde_t* clone_pde(pde_t* cur_dir) {
  pde_t* new_dir;
  pte_t* old_table;
  u32 i;
  u32 heap_sz;
  new_dir = (pde_t*)(get_page());
  memset((u32)new_dir, 0, sizeof(pde_t));
  
  for(i = 0; i < PDE_OFFSET(KERNEL_VIRT_START); ++i) {
    if(cur_dir->tables[i] != 0) {
      old_table = PG_FRAME((u32)cur_dir->tables[i]) + KERNEL_VIRT_OFFSET;
      new_dir->tables[i] = clone_pte(old_table, PG_PRES | PG_USER | PG_READ | PG_COW);
    }
  }

  heap_sz = heap_size();

  // link kernel space, but only for the virtual range 
  for(i = PDE_OFFSET(KERNEL_VIRT_START); i < PDE_OFFSET(HEAP_START + heap_sz) + 1; ++i) {
    if(kernel_dir->tables[i] != 0)
      // change permissions to user, read
      new_dir->tables[i] = PG_FRAME((u32)kernel_dir->tables[i]) | PG_PRES | PG_USER | PG_READ;
  }

  return new_dir;
}

// frees frames, pages, and directory
void free_pde(pde_t* dir) {
  int i, j, do_free, found;
  pte_t* table;
  big_page_desc* big_page_info;
  little_page_desc* little_page_info;

  // first clearn kernel page entries
  for(i = PDE_OFFSET(KERNEL_VIRT_START); i <= 0x3ff; ++i) {
    if(kernel_dir->tables[i] != 0) 
      dir->tables[i] = 0;
  }

  // now free allocated frames, and return page addresses
  for(i = 0; i <= PDE_OFFSET(KERNEL_VIRT_START); ++i) {
    if(dir->tables[i] != 0) {
      table = (pte_t*)(PG_FRAME((u32)dir->tables[i]) + KERNEL_VIRT_OFFSET);
      // assume we can free this table
      do_free = TRUE;
      found = FALSE;

      // its more likely this page is a reference, so we check that first
      for(little_page_info = little_page_desc_list->next; little_page_info != little_page_desc_list && found == FALSE;
	  little_page_info = little_page_info->next) {
	if(little_page_info->addr == (u32)table) {
	  found = TRUE;
	  big_page_info = little_page_info->ref_page;
	  big_page_info->ref_count = big_page_info->ref_count - 1;
	  little_page_info->next->prev = little_page_info->prev;
	  little_page_info->prev->next = little_page_info->next;
	  kfree(little_page_info);

	  if(big_page_info->ref_count == 0)
	    big_page_info->free = TRUE;
	  
	  // if this table is done with, get rid of it
	  if(big_page_info->free == TRUE && big_page_info->in_use == FALSE) {
	    free_pte((pte_t*)big_page_info->addr);
	    big_page_info->next->prev = big_page_info->prev;
	    big_page_info->prev->next = big_page_info->next;
	    kfree((u32)big_page_info);
	  }
	}
      }
    
      // this wasnt a reference to other tables, so see if its being referenced
      if(found == FALSE) {
	for(big_page_info = big_page_desc_list->next; big_page_info != big_page_desc_list;
	    big_page_info = big_page_info->next) {
	  if(big_page_info->addr == (u32)table) {
	    big_page_info->in_use = FALSE;
	    
	    if(big_page_info->ref_count == 0) {
	      big_page_info->next->prev = big_page_info->prev;
	      big_page_info->prev->next = big_page_info->next;
	      kfree((u32)big_page_info);
	    }
	    else
	      do_free = FALSE;
	  }
	}
      }
      
      if(do_free) 
	free_pte(table);
    }
  }
  
  // finally free the directory's address
  free_page(PG_FRAME((u32)dir));
}

// clones a table and marks its frames with the supplied flags
// expercts virt table addr, returns phys table addr
pte_t* clone_pte(pte_t* old_table, u32 flags) {
  int i, j;
  u32* old_frame;
  u32* new_frame;
  u32 found = FALSE;
  big_page_desc* big_page_info;
  little_page_desc* little_page_info;
  pte_t* new_table = (pte_t*)(get_page());
  memset((u32)new_table, 0, sizeof(pte_t));

  // see if the old table is already being referenced
  for(big_page_info = big_page_desc_list->next; big_page_info != big_page_desc_list && found == FALSE;
      big_page_info = big_page_info->next) {
    if(big_page_info->addr == (u32)old_table) {
      found = TRUE;
      big_page_info->ref_count = big_page_info->ref_count + 1;
      little_page_info = (little_page_desc*)kmalloc(sizeof(little_page_desc));
      little_page_info->addr = PG_FRAME((u32)new_table);
      little_page_info->ref_page = big_page_info;
      little_page_info->prev = little_page_desc_list;
      little_page_info->next = little_page_desc_list->next;
      little_page_desc_list->next->prev = little_page_info;
      little_page_desc_list->next = little_page_info;
    }
  }

  // see if the old table references another table, so we dont duplicate refs
  for(little_page_info = little_page_desc_list->next; little_page_info != little_page_desc_list && found == FALSE;
      little_page_info = little_page_info->next) {
    if(little_page_info->addr == (u32)old_table) {
      found = TRUE;
      big_page_info = little_page_info->ref_page;
      big_page_info->ref_count = big_page_info->ref_count + 1;
      little_page_info = (little_page_desc*)kmalloc(sizeof(little_page_desc));
      little_page_info->addr = PG_FRAME((u32)new_table);
      little_page_info->ref_page = big_page_info;
      little_page_info->prev = little_page_desc_list;
      little_page_info->next = little_page_desc_list->next;
      little_page_desc_list->next->prev = little_page_info;
      little_page_desc_list->next = little_page_info;
    }
  }

  // if it wasnt, make it be referenced
  if(!found) {
    big_page_info = (big_page_desc*)kmalloc(sizeof(big_page_desc));
    big_page_info->ref_count = 1;
    big_page_info->addr = PG_FRAME((u32)old_table);
    big_page_info->free = FALSE;
    big_page_info->in_use = TRUE;
    big_page_info->prev = big_page_desc_list;
    big_page_info->next = big_page_desc_list->next;
    big_page_desc_list->next->prev = big_page_info;
    big_page_desc_list->next = big_page_info;
    little_page_info = (little_page_desc*)kmalloc(sizeof(little_page_desc));
    little_page_info->addr = PG_FRAME((u32)new_table);
    little_page_info->ref_page = big_page_info;
    little_page_info->prev = little_page_desc_list;
    little_page_info->next = little_page_desc_list->next;
    little_page_desc_list->next->prev = little_page_info;
    little_page_desc_list->next = little_page_info;
  }     

  for(i = 0; i < 1024; ++i) {
    if(old_table->pages[i] != 0) {
      new_table->pages[i] = (page_t)((u32)old_table->pages[i] | flags);
      //u32 frm = get_frame() | PG_PRES | PG_WRITE;
      //new_table->pages[i] = (page_t)frm;
      //copy_frame((u32)PG_FRAME(old_table->pages[i]), frm);
    }
  }

  return (pte_t*)(((u32)new_table | flags) - KERNEL_VIRT_OFFSET);
}

// expects virt table addr
void free_pte(pte_t* table) {
  int j;

  for(j = 0; j <= 0x3ff; ++j) {
    if(table->pages[j] != 0)
      // if this was marked as a reference to another frame, it is not original so dont free it
      if(!((u32)table->pages[j] & PG_REF))
	free_frame(PG_FRAME((u32)table->pages[j]));
  }

  free_page(table);
}

// temporarily maps both frames in current directory so we can copy
void copy_frame(u32 src, u32 dst) {
  pte_t* table = (get_page() | PG_PRES | PG_WRITE) - KERNEL_VIRT_OFFSET;
  // for now 0xF0000000 is safe because nothing should be there
  cur_dir->tables[PDE_OFFSET(0xF0000000)] = table;
  table = (pte_t*)(PG_FRAME((u32)table) + KERNEL_VIRT_OFFSET);
  table->pages[0] = dst | PG_PRES | PG_WRITE;
  table->pages[1] = src | PG_PRES | PG_WRITE;
  asm volatile ("invlpg 0xF0000000");
  asm volatile ("invlpg 0xF0001000");
  memcpy((u32*)0xF0000000, (u32*)0xF0001000, 0x1000);
  cur_dir->tables[PDE_OFFSET(0xF0000000)] = 0;
  free_page((u32)PG_FRAME((u32)table));
}

void switch_page_directory(pde_t* dir) {
  cur_dir = dir;
  load_CR3((u32)dir - KERNEL_VIRT_OFFSET);
  enable_paging();
}

u32 virt_to_phys(pde_t* dir, u32 addr) {
  u32 result;
  pte_t* table = (pte_t*)PG_FRAME((u32)dir->tables[PDE_OFFSET(addr)]);
  result = PG_FRAME(table->pages[PTE_OFFSET(addr)]);
  return result + (addr & 0x00000FFF);
}

void debug_mm() {
  debug_virt_mm();
  debug_phys_mm();
}

void debug_big_desc_list() {
  big_page_desc* tmp;
  kprintf("big list: ");

  for(tmp = big_page_desc_list->next; tmp != big_page_desc_list; tmp = tmp->next)
    kprintf("addr: %x ref_count: %d in_use: %d  \n          ", tmp->addr, tmp->ref_count, tmp->in_use);
  kprintf("\n");
}

void debug_little_desc_list() {
  little_page_desc* tmp;
  kprintf("little list: ");

  for(tmp = little_page_desc_list->next; tmp != little_page_desc_list; tmp = tmp->next)
    kprintf("addr: %x ref_page: %x\n             ", tmp->addr, tmp->ref_page->addr);
  kprintf("\n");
}

void debug_list() {
 little_page_desc* tmp;
  kprintf("list: ");

  for(tmp = little_page_desc_list->next; tmp != little_page_desc_list; tmp = tmp->next)
    kprintf("addr: %x next: %x prev: %x\n", tmp->addr, tmp->next, tmp->prev);
  kprintf("\n");
}


 
